# coding = utf-8


'''
加载数据的代码
'''
from torch.utils import data
import os
import numpy as np
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from skimage.transform import resize

class IRCAD_DATA(data.Dataset):
    def __init__(self, origion_data_path, hdense_data_path, input_size, input_cols, mean):

        self._origion_data_path = origion_data_path
        self._hdense_data_path = hdense_data_path
        self._input_size = input_size
        self._input_cols = input_cols
        self._mean = mean


        self._get_data(train_file="/datasets/3Dircadb/chengkung/train.txt", valid_file="/datasets/3Dircadb/chengkung/val.txt")
        self._split_subset()


    def _get_data(self, train_file, valid_file):
        def read_txt(file):
            d = []
            f = open(file, 'r')
            for line in f:
                d.append(int(line))
            return d

        self._train_case = read_txt(train_file)
        self._valid_case = read_txt(valid_file)

        self._case_id = self._train_case + self._valid_case

        self._indices = list(range(len(self._case_id)))
        self._train_indices = self._indices[:len(self._train_case)]
        self._valid_indices = self._indices[len(self._train_case):]

        self._img_list = []
        self._tumor_list = []
        self._tumorline_list = []
        self._liverline_list = []
        self._liverBox_list = []

        for case_id in self._case_id:
            self._img_list.append(os.path.join(self._origion_data_path, "case_{}/image.npy".format(str(case_id).zfill(5))))
            self._tumor_list.append(os.path.join(self._origion_data_path, "case_{}/segment.npy".format(str(case_id).zfill(5))))
            self._tumorline_list.append(os.path.join(self._hdense_data_path, "TumorPixels/case_{}.txt".format(str(case_id).zfill(5))))
            self._liverline_list.append(os.path.join(self._hdense_data_path, "LiverPixels/case_{}.txt".format(str(case_id).zfill(5))))
            self._liverBox_list.append(os.path.join(self._hdense_data_path, "LiverBox/case_{}.txt".format(str(case_id).zfill(5))))

    def _split_subset(self):
            self._train_dataset = data.Subset(self, self._train_indices)
            self._valid_dataset = data.Subset(self, self._valid_indices)


    def __len__(self):
        return len(self._indices)

    def __getitem__(self, i):
        image_file_name = self._img_list[i]
        mask_file_name = self._tumor_list[i]
        case_id = self._case_id[i]

        image =  np.load(image_file_name)
        mask = np.load(mask_file_name)

        image = image.transpose([1,2,0])
        mask = mask.transpose([1,2,0])

        maxmin = np.loadtxt(self._liverBox_list[i], delimiter=' ')
        minindex = maxmin[0:3]
        maxindex = maxmin[3:6]
        minindex = np.array(minindex, dtype='int')
        maxindex = np.array(maxindex, dtype='int')
        minindex[0] = max(minindex[0] - 3, 0)
        minindex[1] = max(minindex[1] - 3, 0)
        minindex[2] = max(minindex[2] - 3, 0)
        maxindex[0] = min(image.shape[0], maxindex[0] + 3)
        maxindex[1] = min(image.shape[1], maxindex[1] + 3)
        maxindex[2] = min(image.shape[2], maxindex[2] + 3)

        f1 = open(self._tumorline_list[i], 'r')
        tumorline = f1.readlines()
        f1.close()

        f2 = open(self._liverline_list[i], 'r')
        liverline = f2.readlines()
        f2.close()

        num = np.random.randint(0, 6)
        if num < 3 and len(tumorline) > 1:
            lines = tumorline
            numid = len(tumorline)
        else:
            lines = liverline
            numid = len(liverline)
        Parameter_List = [image, mask, lines, numid, minindex, maxindex]

        while True:
           X, Y = self.load_seq_crop_data_masktumor_try(Parameter_List)
           c0 = np.sum(Y == 0)
           c1 = np.sum(Y == 1)
           c2 = np.sum(Y == 2)
           break
        #print(c0, ' ', c1, ' ', c2)
        # if (c0 * c1 * c2 == 0):
        #   continue
        return (X.transpose([2, 0, 1])).astype(np.float32), (Y.transpose([2, 0, 1])).astype(np.int)

    def load_seq_crop_data_masktumor_try(self, Parameter_List):
        img = Parameter_List[0]
        tumor = Parameter_List[1]
        lines = Parameter_List[2]
        numid = Parameter_List[3]
        minindex = Parameter_List[4]
        maxindex = Parameter_List[5]
        #  randomly scale
        scale = np.random.uniform(0.8, 1.2)
        deps = int(self._input_size * scale)
        rows = int(self._input_size * scale)
        cols = self._input_cols

        sed = np.random.randint(1, numid)
        cen = lines[sed - 1]
        cen = np.fromstring(cen, dtype=int, sep=' ')
        # print (cen)
        a = min(max(minindex[0] + deps / 2, cen[0]), maxindex[0] - deps / 2 - 1)
        b = min(max(minindex[1] + rows / 2, cen[1]), maxindex[1] - rows / 2 - 1)
        c = min(max(minindex[2] + cols / 2, cen[2]), maxindex[2] - cols / 2 - 1)
        # print(a - deps // 2, ' ', c - args.input_cols // 2, ' ', c + args.input_cols // 2)
        cropp_img = img[int(a - deps / 2):int(a + deps / 2), int(b - rows / 2):int(b + rows / 2),
                    int(c - self._input_cols / 2): int(c + self._input_cols / 2)].copy()

        cropp_tumor = tumor[int(a - deps / 2):int(a + deps / 2), int(b - rows / 2):int(b + rows / 2),
                      int(c - self._input_cols / 2):int(c + self._input_cols / 2)].copy()

        cropp_img -= self._mean
        # randomly flipping
        flip_num = np.random.randint(0, 8)
        if flip_num == 1:
            cropp_img = np.flipud(cropp_img)
            cropp_tumor = np.flipud(cropp_tumor)
        elif flip_num == 2:
            cropp_img = np.fliplr(cropp_img)
            cropp_tumor = np.fliplr(cropp_tumor)
        elif flip_num == 3:
            cropp_img = np.rot90(cropp_img, k=1, axes=(1, 0))
            cropp_tumor = np.rot90(cropp_tumor, k=1, axes=(1, 0))
        elif flip_num == 4:
            cropp_img = np.rot90(cropp_img, k=3, axes=(1, 0))
            cropp_tumor = np.rot90(cropp_tumor, k=3, axes=(1, 0))
        elif flip_num == 5:
            cropp_img = np.fliplr(cropp_img)
            cropp_tumor = np.fliplr(cropp_tumor)
            cropp_img = np.rot90(cropp_img, k=1, axes=(1, 0))
            cropp_tumor = np.rot90(cropp_tumor, k=1, axes=(1, 0))
        elif flip_num == 6:
            cropp_img = np.fliplr(cropp_img)
            cropp_tumor = np.fliplr(cropp_tumor)
            cropp_img = np.rot90(cropp_img, k=3, axes=(1, 0))
            cropp_tumor = np.rot90(cropp_tumor, k=3, axes=(1, 0))
        elif flip_num == 7:
            cropp_img = np.flipud(cropp_img)
            cropp_tumor = np.flipud(cropp_tumor)
            cropp_img = np.fliplr(cropp_img)
            cropp_tumor = np.fliplr(cropp_tumor)
        #
        cropp_tumor = resize(cropp_tumor, (self._input_size, self._input_size, self._input_cols), order=0, mode='edge',
                             cval=0, clip=True, preserve_range=True)
        cropp_img = resize(cropp_img, (self._input_size, self._input_size, self._input_cols), order=3, mode='constant',
                           cval=0, clip=True, preserve_range=True)
        return cropp_img, cropp_tumor

    @property
    def valid_dataset(self):
        return self._valid_dataset

    @property
    def train_dataset(self):
        return self._train_dataset

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    ircad = IRCAD_DATA(origion_data_path="/datasets/3Dircadb/origion", hdense_data_path="/datasets/3Dircadb/HDenseNet",
                       input_size=224, input_cols=8, mean=0)

    subset = ircad.valid_dataset

    sampler = SequentialSampler(subset)
    data_loader = DataLoader(subset, batch_size=1)

    for index,(data,mask) in enumerate(data_loader):
        print(data.shape)
        print(mask.shape)
        data = data.squeeze().numpy()
        mask = mask.squeeze().numpy()
        for i in range(data.shape[0]):
           plt.subplot(1, 2, 1)
           plt.imshow(data[i], cmap="gray")
           plt.subplot(1, 2, 2)
           plt.imshow(mask[i], cmap="gray")
           plt.show()

        print()